package edu.washington.ling.syntric.eval;

import java.io.IOException;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.stanford.nlp.util.Pair;
import edu.washington.ling.syntric.io.Example;
import edu.washington.ling.syntric.io.ExampleDoc;
import edu.washington.ling.syntric.io.ExampleDocReader;
import edu.washington.ling.syntric.io.TreeReader;
import edu.washington.ling.syntric.io.TreeRepresentation;
import edu.washington.ling.syntric.struct.error.InvalidValueException;
import edu.washington.ling.syntric.struct.error.PrintingExceptionHandler;
import edu.washington.ling.syntric.struct.error.UnreadableRepresentationException;
import edu.washington.ling.syntric.struct.tree.Node;
import edu.washington.ling.syntric.struct.tree.Token;
import edu.washington.ling.syntric.struct.tree.Tree;
import edu.washington.ling.syntric.util.ConvertUtils;


class TreeSpan {
	protected int left;
	protected int right;
	
	public TreeSpan(int left, int right) {
		this.left = left;
		this.right = right;
	}
	public int getLeft() {
		return left;
	}
	public int getRight() {
		return right;
	}
	@Override
	public boolean equals(Object other) {
		boolean isEqual = false;
		if (other instanceof TreeSpan) {
			TreeSpan otherTreeSpan = (TreeSpan) other;
			if (this.left == otherTreeSpan.left && this.right == otherTreeSpan.right) {
				isEqual = true;
			}	
		}
		return isEqual;
	}
	@Override
	public int hashCode() {
		int hashCode = new Pair<Integer, Integer>(left, right).hashCode();
		return hashCode;
	}
}

class LabeledTreeSpan extends TreeSpan {
	private String label;
	
	public LabeledTreeSpan(int left, int right, String label) {
		super(left, right);
		this.label = label;
	}
	@Override
	public boolean equals(Object other) {
		boolean isEqual = false;
		if (other instanceof TreeSpan) {
			TreeSpan otherTreeSpan = (TreeSpan) other;
			if (this.left == otherTreeSpan.left && this.right == otherTreeSpan.right
					&& this.label.equals(label)) {
				isEqual = true;
			}	
		}
		return isEqual;
	}
	@Override
	public int hashCode() {
		int hashCode = new Pair<Integer, Integer>(left, right).hashCode() + label.hashCode();
		return hashCode;
	}
	
}

class EvalTracker {
	int tpCount;
	int fpCount;
	int fnCount;
	
	public EvalTracker() {
		tpCount = 0;
		fpCount = 0;
		fnCount = 0;
	}
	
	public void addTruePos(int count) {
		tpCount += count;
	}
	public void addFalsePos(int count) {
		fpCount += count;
	}
	public void addFalseNeg(int count) {
		fnCount += count;
	}
	public void addStats(EvalTracker tracker) {
		tpCount += tracker.tpCount;
		fpCount += tracker.fpCount;
		fnCount += tracker.fnCount;
	}
	public int truePos() {
		return tpCount;
	}
	public int falsePos() {
		return fpCount;
	}
	public int falseNeg() {
		return fnCount;
	}
	public double precision() {
		return ((double) tpCount / (tpCount + fpCount));
	}
	public double recall() {
		return ((double) tpCount / (tpCount + fnCount));
	}
	public double f1() {
		double f1 = 0.0;
		double recall = recall();
		double precision = precision();
		if (recall + precision > 0) {
			f1 = 2*recall*precision / (recall + precision);
		}
		return f1;
	}
	
}

public class Evaluate {
	public enum EvalLabels {
		NoLabels, SimpleLabels, ExtendedLabels
	}
	public enum EvalMethod {
		MaxEval, FirstEval
	}
	private static EvaluateOptions opts = new EvaluateOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	private static Node getLeftEdgeNode(Node n) {
		int numChildren = n.getPSChildren().size();
		if (numChildren > 0) {
			return getLeftEdgeNode(n.getPSChildren().get(0));
		}
		else {
			return n;
		}
	}
	private static Node getRightEdgeNode(Node n) {
		int numChildren = n.getPSChildren().size();
		if (numChildren > 0) {
			return getRightEdgeNode(n.getPSChildren().get(numChildren - 1));
		}
		else {
			return n;
		}
	}
	private static void getSpans(Node n, Tree tree, EvalLabels evalLabels, List<TreeSpan> treeSpans) {
		if (!(n instanceof Token)) {
			Node leftNode = getLeftEdgeNode(n);
			Node rightNode = getRightEdgeNode(n);
			if (leftNode instanceof Token && rightNode instanceof Token) {
				int left = tree.getTokenPosition((Token) leftNode);
				int right = tree.getTokenPosition((Token) rightNode);
				TreeSpan span;
				if (evalLabels.equals(EvalLabels.ExtendedLabels)) {
					String label = n.getLabel();
					span = new LabeledTreeSpan(left, right, label);
				} else if (evalLabels.equals(EvalLabels.SimpleLabels)) {
					String label = n.getLabel().split("-")[0];
					span = new LabeledTreeSpan(left, right, label);
				} else {
					span = new TreeSpan(left, right);
				}
				treeSpans.add(span);
				if (left != right) {
					for (Node child: n.getPSChildren()) {
						getSpans(child, tree, evalLabels, treeSpans);
					}
				}
			}
		}
	}
	private static Set<TreeSpan> getSpans(Tree tree, EvalLabels evalLabels) {
		List<TreeSpan> spans = new LinkedList<TreeSpan>();
		getSpans(tree.getPSRoot(), tree, evalLabels, spans);
		Set<TreeSpan> spanSet = new HashSet<TreeSpan>(spans);
		return spanSet;
	}

	private static EvalTracker evalTree(Tree systemTree, Tree goldTree, EvalLabels evalLabels) {
		EvalTracker evalTracker = new EvalTracker();
		Set<TreeSpan> systemSpans = getSpans(systemTree, evalLabels);
		Set<TreeSpan> goldSpans = getSpans(goldTree, evalLabels);
		Set<TreeSpan> sharedSpans = new HashSet<TreeSpan>(systemSpans);
		sharedSpans.retainAll(goldSpans);
		evalTracker.addTruePos(sharedSpans.size());
		evalTracker.addFalsePos(systemSpans.size() - sharedSpans.size());
		evalTracker.addFalseNeg(goldSpans.size() - sharedSpans.size());
		return evalTracker;
	}
	private static EvalTracker evalMissingTree(Tree goldTree, EvalLabels evalLabels) {
		EvalTracker evalTracker = new EvalTracker();
		Set<TreeSpan> goldSpans = getSpans(goldTree, evalLabels);
		evalTracker.addFalseNeg(goldSpans.size());
		return evalTracker;
	}
	public static List<String> getCompatibleExampleIDs(ExampleDoc exf) {
		List<String> compatExampleIDs = new LinkedList<String>();
		for (String exampleID: exf.getIDs()) {
			if (ConvertUtils.isCompatible(exf.getExampleWithID(exampleID))) {
				compatExampleIDs.add(exampleID);
			}
		}
		return compatExampleIDs;
	}
	private static EvalTracker maxEval(List<EvalTracker> evalTrackers) {
		double maxF1 = 0.0;
		EvalTracker maxTracker = new EvalTracker();
		for (EvalTracker evalTracker: evalTrackers) {
			double thisF1 = evalTracker.f1();
			if (thisF1 > maxF1) {
				maxF1 = thisF1;
				maxTracker = evalTracker;
			}
		}
		return maxTracker;
	}
	public static void evaluate(ExampleDoc exf, List<String> exampleIDs) {
		int numExamples = 0;
		int numWithOutput = 0;
		EvalTracker trackerRestricted = new EvalTracker();
		EvalTracker tracker = new EvalTracker();
		for (String exampleID: exampleIDs) {
			Example exr = exf.getExampleWithID(exampleID);
			if (exr.hasTreeRep("PS")) {
				numExamples++;
				try {
					TreeRepresentation goldTreeRep = exr.getTreeRep("PS");
					TreeReader goldTreeReader = TreeReader.newReader(goldTreeRep.getFormat());
					Tree goldTree = goldTreeReader.readTree(goldTreeRep);
					if (exr.hasSystemTreeReps()) {
						numWithOutput++;
						List<TreeRepresentation> systemTreeReps = exr.getSystemTreeReps();
						List<EvalTracker> systemTreeTrackers = new LinkedList<EvalTracker>();
						for (TreeRepresentation systemTreeRep: systemTreeReps) {
							TreeReader systemTreeReader = TreeReader.newReader(systemTreeRep.getFormat());
							Tree systemTree = systemTreeReader.readTree(systemTreeRep);
							systemTreeTrackers.add(evalTree(systemTree, goldTree, opts.evalLabels));
						}
						EvalTracker systemTreeTracker;
						if (opts.evalMethod.equals(EvalMethod.MaxEval)) {
							systemTreeTracker = maxEval(systemTreeTrackers);
						} else if (opts.evalMethod.equals(EvalMethod.FirstEval)){
							systemTreeTracker = systemTreeTrackers.get(0);
						} else {
						systemTreeTracker = systemTreeTrackers.get(0);
						}
						System.out.println(String.format("%-25s | TP %4d | FP %4d | FN %4d | P %3.2f | R %3.2f | F1 %3.2f |",
								exampleID, systemTreeTracker.truePos(), systemTreeTracker.falsePos(), systemTreeTracker.falseNeg(),
								systemTreeTracker.precision(), systemTreeTracker.recall(), systemTreeTracker.f1()));
						trackerRestricted.addStats(systemTreeTracker);
						tracker.addStats(systemTreeTracker);
					} else {
						EvalTracker missingTreeTracker = evalMissingTree(goldTree, opts.evalLabels);
						tracker.addStats(missingTreeTracker);
					}
				}
				catch (UnreadableRepresentationException e) {
					String errStr = String.format("Could not process representation for example '%s'", exampleID);
					pHandler.handle(e, errStr);
				}
				catch (InvalidValueException e) {
					String errStr = String.format("Could not process representation for example '%s'", exampleID);
					pHandler.handle(e, errStr);
				}
			}
		}
		System.out.println("Evaluation was successful");
		System.out.println();
		System.out.println(String.format("Settings:'%s','%s'",opts.evalLabels, opts.evalMethod));
		System.out.println();
		System.out.println(String.format("Results restricted to %d examples with tree output:", numWithOutput));
		System.out.println(String.format("| TP %4d | FP %4d | FN %4d | P %3.2f | R %3.2f | F1 %3.2f |",
				trackerRestricted.truePos(), trackerRestricted.falsePos(), trackerRestricted.falseNeg(),
				trackerRestricted.precision(), trackerRestricted.recall(), trackerRestricted.f1()));
		System.out.println();
		System.out.println(String.format("Results for all %d examples:", numExamples));
		System.out.println(String.format("| TP %4d | FP %4d | FN %4d | P %3.2f | R %3.2f | F1 %3.2f |",
				tracker.truePos(), tracker.falsePos(), tracker.falseNeg(),
				tracker.precision(), tracker.recall(), tracker.f1()));
		
	}
	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			List<String> exampleIDs;
			if (opts.compatOnly) {
				exampleIDs = getCompatibleExampleIDs(exf);
			} else {
				exampleIDs = new LinkedList<String>(exf.getIDs());
			}
			evaluate(exf, exampleIDs);
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("Usage: Evaluate");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(-1);
		}
		catch(IOException e) {
			String errStr = String.format("Example file '%s' could not be read", opts.inputFilePath);
			pHandler.handle(e, errStr, System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}
}
